package org.brakepedal.server.util;

import org.jetbrains.annotations.NotNull;

import java.lang.reflect.Array;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.*;

/**
 * Date: 2023/10/12 15:05
 *

 */
public class Matrix<ElementType> {
    private final ElementType[][] matrix;

    private final AtomicReference<MatrixOperator<ElementType>> operator = new AtomicReference<>();

    // alise for rowLength, matrix.length
    public final int length;

    public final int columnLength;

    public final int rowLength;

    public Matrix(@NotNull ElementType[] @NotNull [] matrix) {
        check(matrix, () -> new IllegalArgumentException(String.format("'%s' is not a matrix", toString(matrix))));
        this.matrix = matrix;
        length = matrix.length;
        columnLength = matrix[0].length;
        rowLength = matrix.length;
    }

    @FunctionalInterface
    public interface ThreeArgsCellReplacer<CellType> {
        CellType replace(CellType value, Integer rowIndex, Integer columnIndex);
    }

    @FunctionalInterface
    public interface ThreeArgsCellAccessor<CellType> {
        void access(CellType value, Integer rowIndex, Integer columnIndex);
    }

    public static class MatrixOperator<ElementType> {
        private final Matrix<ElementType> matrix;

        public class RowAccessor {
            private final int index;

            public RowAccessor(int index) {
                this.index = index;
            }

            public RowAccessor replace(Function<ElementType, ElementType> replacer) {
                return replace((el, ignore) -> replacer.apply(el));
            }

            public RowAccessor replace(BiFunction<ElementType, Integer, ElementType> replacer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat[index].length; i++) {
                    mat[index][i] = replacer.apply(mat[index][i], i);
                }
                return this;
            }

            public ElementType[] slice() {
                return matrix.get()[index];
            }

            @SuppressWarnings("all")
            public <T> T[] mapTo(Class<T> target, BiFunction<ElementType, Integer, T> mapper) {
                ElementType[] row = matrix.get()[index];
                T[] result = ((T[]) Array.newInstance(target, row.length));
                for (int i = 0; i < row.length; i++) {
                    result[i] = mapper.apply(row[i], i);
                }
                return result;
            }

            public <T> T[] mapTo(Class<T> target, Function<ElementType, T> mapper) {
                return mapTo(target, (el, ignore) -> mapper.apply(el));
            }

            public ElementType column(int columnIndex) {
                return matrix.get()[index][columnIndex];
            }

            public Matrix<ElementType> matrix() {
                return matrix;
            }

            public RowAccessor forEach(BiConsumer<ElementType, Integer> consumer) {
                ElementType[] row = matrix.get()[index];
                for (int i = 0; i < row.length; i++) {
                    consumer.accept(row[i], i);
                }
                return this;
            }

            public RowAccessor forEach(Consumer<ElementType> consumer) {
                return forEach((el, index) -> consumer.accept(el));
            }
        }

        public class ColumnAccessor {
            private final int index;

            public ColumnAccessor(int index) {
                this.index = index;
            }

            public Matrix<ElementType> replace(Function<ElementType, ElementType> replacer) {
                return replace((el, ignore) -> replacer.apply(el));
            }

            public Matrix<ElementType> replace(BiFunction<ElementType, Integer, ElementType> replacer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat.length; i++) {
                    mat[i][index] = replacer.apply(mat[i][index], i);
                }
                return matrix;
            }

            @SuppressWarnings("all")
            public ElementType[] slice() {
                return Arrays.stream(matrix.get()).map(row -> row[index]).toArray(num -> ((ElementType[]) Array.newInstance(matrix.get().getClass().getComponentType().getComponentType(), num)));
            }

            @SuppressWarnings("all")
            public <T> T[] mapTo(Class<T> target, BiFunction<ElementType, Integer, T> mapper) {
                ElementType[][] mat = matrix.get();
                T[] result = ((T[]) Array.newInstance(target, mat[0].length));
                for (int i = 0; i < result.length; i++) {
                    result[i] = mapper.apply(mat[i][index], i);
                }
                return result;
            }

            public <T> T[] mapTo(Class<T> target, Function<ElementType, T> mapper) {
                return mapTo(target, (el, ignore) -> mapper.apply(el));
            }

            public ElementType row(int rowIndex) {
                return matrix.get()[rowIndex][index];
            }

            public Matrix<ElementType> matrix() {
                return matrix;
            }

            public ColumnAccessor forEach(BiConsumer<ElementType, Integer> consumer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat[0].length; i++) {
                    consumer.accept(mat[i][index], i);
                }
                return this;
            }

            public ColumnAccessor forEach(Consumer<ElementType> consumer) {
                return forEach((el, index) -> consumer.accept(el));
            }
        }

        public class RowsAccessor {
            public RowsAccessor forEach(BiConsumer<ElementType[], Integer> consumer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat.length; i++) {
                    consumer.accept(row(i).slice(), i);
                }
                return this;
            }

            public RowsAccessor forEach(Consumer<ElementType[]> consumer) {
                return forEach((el, ignore) -> consumer.accept(el));
            }

            public RowsAccessor forEach(ThreeArgsCellAccessor<ElementType> accessor) {
                return forEach((row, rowIndex) -> {
                    for (int columnIndex = 0; columnIndex < row.length; columnIndex++) {
                        accessor.access(row[columnIndex], rowIndex, columnIndex);
                    }
                });
            }

            public RowsAccessor replace(ThreeArgsCellReplacer<ElementType> replacer) {
                return forEach((row, rowIndex) -> {
                    for (int columnIndex = 0; columnIndex < row.length; columnIndex++) {
                        row[columnIndex] = replacer.replace(row[columnIndex], rowIndex, columnIndex);
                    }
                });
            }

            /**
             * @param mapper {@code (T[] row, Integer rowIndex)-> T[]}
             */
            public <T> T[] mapTo(Class<T> target, BiFunction<ElementType[], Integer, T> mapper) {
                @SuppressWarnings("unchecked")
                T[] result = ((T[]) Array.newInstance(target, matrix.get().length));
                for (int i = 0; i < result.length; i++) {
                    result[i] = mapper.apply(row(i).slice(), i);
                }
                return result;
            }

            /**
             * @param mapper {@code (T[] row, Integer rowIndex)-> T[]}
             */
            public <T> T[] mapTo(IntFunction<T[]> target, BiFunction<ElementType[], Integer, T> mapper) {
                T[] result = target.apply(matrix.rowLength);
                if(result.length != matrix.rowLength) {
                    result = Arrays.copyOf(result, matrix.rowLength);
                }
                for (int i = 0; i < result.length; i++) {
                    result[i] = mapper.apply(row(i).slice(), i);
                }
                return result;
            }

            public <T> T[] mapTo(Class<T> target, Function<ElementType[], T> mapper) {
                return mapTo(target, (el, ignore) -> mapper.apply(el));
            }

            public <T> T[] mapTo(IntFunction<T[]> target, Function<ElementType[], T> mapper) {
                return mapTo(target, (el, ignore) -> mapper.apply(el));
            }

            public Matrix<ElementType> matrix() {
                return matrix;
            }
        }

        public class ColumnsAccessor {
            public ColumnsAccessor forEach(BiConsumer<ElementType[], Integer> consumer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat[0].length; i++) {
                    consumer.accept(column(i).slice(), i);
                }
                return this;
            }

            public ColumnsAccessor forEach(Consumer<ElementType[]> consumer) {
                return forEach((el, ignore) -> consumer.accept(el));
            }

            public ColumnsAccessor forEach(Runnable runnable) {
                return forEach((el, ignore) -> runnable.run());
            }

            public <T> T[] mapTo(Class<T> target, BiFunction<ElementType[], Integer, T> mapper) {
                @SuppressWarnings("unchecked")
                T[] result = (T[]) Array.newInstance(target, matrix.columnLength);
                for (int i = 0; i < result.length; i++) {
                    ElementType[] column = column(i).slice();
                    T apply = mapper.apply(column, i);
                    result[i] = apply;
                }
                return result;
            }

            @SuppressWarnings("all")
            /*public <T> T[] map(BiFunction<ElementType[], Integer, T> mapper) {
                T[] result = (T[]) Array.newInstance(matrix.get()[0][0].getClass(), matrix.get()[0].length);
                try {
                    GenericTypeResolver.resolveReturnType(getClass().getMethod("map", BiFunction.class), getClass());
                } catch (NoSuchMethodException e) {
                    throw new RuntimeException(e);
                }
                for (int i = 0; i < result.length; i++) {
                    ElementType[] column = column(i).slice();
                    System.out.println("get column result: "+column);
                    T apply = mapper.apply((ElementType[])column, i);
                    System.out.println("get apply result: "+apply);
                    result[i] = apply;
                }
                return result;
            }*/

            public <T> T[] mapTo(Class<T> target, Function<ElementType[], T> mapper) {
                return mapTo(target, (el, ignore) -> mapper.apply(el));
            }

            public ColumnsAccessor replace(BiFunction<ElementType[], Integer, ElementType[]> replacer) {
                ElementType[][] mat = matrix.get();
                for (int i = 0; i < mat.length; i++) {
                    mat[i] = replacer.apply(row(i).slice(), i);
                }
                return this;
            }

            public ColumnsAccessor replace(Function<ElementType[], ElementType[]> replacer) {
                return replace((el, ignore) -> replacer.apply(el));
            }

            public Matrix<ElementType> matrix() {
                return matrix;
            }
        }

        public MatrixOperator(Matrix<ElementType> matrix) {
            this.matrix = matrix;
        }

        /**
         * @param index 行索引
         * @return 遍历改行的访问器
         */
        public RowAccessor row(int index) {
            return new RowAccessor(index);
        }

        /**
         * @param index 列索引
         * @return 遍历改列的访问器
         */
        public ColumnAccessor column(int index) {
            return new ColumnAccessor(index);
        }

        /**
         * 以行遍历的形式遍历矩阵
         *
         * @return {@code RowsAccessor} 行遍历器
         */
        public RowsAccessor rows() {
            return new RowsAccessor();
        }

        /**
         * 以列遍历的形式遍历矩阵
         *
         * @return {@code ColumnsAccessor} 列遍历器
         */
        public ColumnsAccessor columns() {
            return new ColumnsAccessor();
        }
    }

    /**
     * @param matrix 二位数组
     * @return 是否是矩阵
     */
    public static boolean check(Object[][] matrix) {
        try {
            int rowLength = matrix[0].length;
            for (int i = 1; i < matrix.length; i++) {
                if (matrix[i].length != rowLength) {
                    return false;
                }
            }
        } catch (ArrayIndexOutOfBoundsException e) {
            return false;
        }
        return true;
    }

    public static void check(Object[][] matrix, Supplier<RuntimeException> exceptionSupplier) {
        if (!check(matrix)) {
            throw exceptionSupplier.get();
        }
    }

    public static <T> Matrix<T> of(T[][] matrix) {
        return new Matrix<>(matrix);
    }

    public static <T> Matrix<T> of(T[][] matrix, Supplier<RuntimeException> exceptionSupplier) {
        try {
            return of(matrix);
        } catch (RuntimeException e) {
            throw exceptionSupplier.get();
        }
    }

    public static <T> Matrix<T> of(T[][] matrix, Function<RuntimeException, RuntimeException> exceptionAdaptor) {
        try {
            return of(matrix);
        } catch (RuntimeException e) {
            throw exceptionAdaptor.apply(e);
        }
    }

    public ElementType[][] get() {
        return matrix;
    }

    public boolean isSquare() {
        final int order = matrix.length;
        for (ElementType[] row : matrix) {
            if (row.length != order) {
                return false;
            }
        }
        return true;
    }

    /**
     * 使用涉及行列的操作
     */
    public MatrixOperator<ElementType> ops() {
        operator.getAndUpdate(oldValue->{
            if(oldValue == null) {
                return new MatrixOperator<>(this);
            }
            return oldValue;
        });
        return operator.get();
    }

    @Override
    public String toString() {
        return toString(matrix);
    }

    public static <T> String toString(T[][] matrix) {
        return Arrays.deepToString(matrix).replaceAll("],", "],\n");
    }
}
